import torch
import torch.nn as nn
import torch.nn.functional as F


class block(nn.Module):
    def __init__(self, in_channel):
        super().__init__()
        self.part1 = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=24, kernel_size=1, stride=1, padding=0),
            nn.ReLU()
        )
        self.part2 = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=16, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=16, out_channels=24, kernel_size=3, stride=1, padding=1),
            nn.ReLU()
        )
        self.part3 = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=16, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=in_channel, out_channels=24, kernel_size=5, stride=1, padding=2),
            nn.ReLU()
        )
        self.part4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=in_channel, out_channels=24, kernel_size=1, stride=1, padding=0),
            nn.ReLU()
        )

    def forward(self, x):
        x1 = self.part1(x)
        x2 = self.part2(x)
        x3 = self.part3(x)
        x4 = self.part4(x)
        return torch.cat((x1, x2, x3, x4), dim=1)


class inception(nn.Module):
    def __init__(self):
        super(inception, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5, stride=1, padding=0)
        self.incep1 = block(in_channel=10)
        self.conv2 = nn.Conv2d(in_channels=96, out_channels=20, kernel_size=5, stride=1, padding=0)
        self.incep2 = block(in_channel=20)
        self.maxPool = nn.MaxPool2d(2)
        self.rule = nn.ReLU()
        self.fc = nn.Linear(in_features=1536, out_features=10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.rule(x)
        x = self.maxPool(x)
        x = self.incep1(x)
        x = self.conv2(x)
        x = self.rule(x)
        x = self.maxPool(x)
        x = self.incep2(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        output = F.log_softmax(x, dim=1)
        return output
    # TODO





